/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <cmath>
#include <iostream>
#include "parvec_axpby.h"
#include "parvec_dot.h"
#include "par_cg_solver.h"

namespace YHAMG
{

Par_CG_Solver::Par_CG_Solver(const Par_Operator& _A)
	: A(&_A),
	P(0),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_CG_Solver::Par_CG_Solver(const Par_Operator& _A, const Par_Operator& _P)
	: A(&_A),
	P(&_P),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_CG_Solver::Par_CG_Solver(const Par_Operator& _A, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(0),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}

Par_CG_Solver::Par_CG_Solver(const Par_Operator& _A, const Par_Operator& _P, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(&_P),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}

void Par_CG_Solver::operator()(const Par_Vector& b, const Par_Vector& x, int& iter, double& relres) const
{
	MPI_Comm comm = A->comm;

	int rank;
	MPI_Comm_rank(comm, &rank);

	int n = A->OutSize();

	double normb, res, alpha, beta, rho, rho1;
	Par_Vector r(comm), z(comm), p(comm), v(comm);

	r.Resize(n);
	z.Resize(n);
	p.Resize(n);
	v.Resize(n);

	if (PrintStats && rank == 0)
	{
		std::cout << "Iter\tResidual\n";
		std::cout << "----------------------------------------------------------------\n";
	}

	iter = 0;

	A->Apply(x, r);
	ParVec_AXPBY(-1.0, r, 1.0, b, r);

	res = sqrt(ParVec_Dot(r, r));
	normb = sqrt(ParVec_Dot(b, b));

	while (1)
	{
		if (PrintStats && rank == 0)
			std::cout << iter << "\t" << res << "\n";

		if (res / normb <= Tolerance || iter == MaxIters) break;

		if (P)
			P->Apply(r, z);
		else
			z.Copy(r);

		rho = ParVec_Dot(r, z);

		if (iter == 0)
			p.Copy(z);
		else
		{
			beta = rho / rho1;
			ParVec_AXPBY(1.0, z, beta, p, p);
		}

		A->Apply(p, v);

		alpha = rho / ParVec_Dot(p, v);
		rho1 = rho;

		x.AddScaled(alpha, p);
		r.AddScaled(-alpha, v);

		res = sqrt(ParVec_Dot(r, r));

		++iter;
	}

	relres = res / normb;

	if (PrintStats && rank == 0)
	{
		std::cout << "----------------------------------------------------------------\n";
		std::cout << "Iterations: " << iter << "\n";
		std::cout << "Final Relative Residual: " << relres << "\n";
	}
}

}